{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "tZnIXBfrRpex"
   },
   "source": [
    "# Question Answering with Langchain and OpenAI\n",
    "\n",
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/elastic/elasticsearch-labs/blob/main/notebooks/generative-ai/question-answering.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
    "\n",
    "This interactive notebook uses Langchain to split fictional workplace documents into passages and uses OpenAI to transform these passages into embeddings and store them into Elasticsearch.\n",
    "\n",
    "\n",
    "![image.png]()\n",
    "\n",
    "Then when we ask a question, we retrieve the relevant passages from the vector store and use langchain and OpenAI to provide a summary for the question."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "GyAst2W-VpHb"
   },
   "source": [
    "## Install required packages\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "33A-cP-XvFCr"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "langserve 0.0.21 requires pydantic<2,>=1, but you have pydantic 2.3.0 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.3.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!python3 -m pip install -qU langchain langchain-elasticsearch openai==0.28.1 tiktoken jq"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "qtEOCsCLWCZp"
   },
   "source": [
    "## Connect to Elasticsearch\n",
    "\n",
    "ℹ️ We're using an Elastic Cloud deployment of Elasticsearch for this notebook. If you don't have an Elastic Cloud deployment, sign up [here](https://cloud.elastic.co/registration?utm_source=github&utm_content=elasticsearch-labs-notebook) for a free trial. \n",
    "\n",
    "We'll use the **Cloud ID** to identify our deployment, because we are using Elastic Cloud deployment. To find the Cloud ID for your deployment, go to https://cloud.elastic.co/deployments and select your deployment.\n",
    "\n",
    "\n",
    "We will use [ElasticsearchStore](https://api.python.langchain.com/en/latest/vectorstores/langchain.vectorstores.elasticsearch.ElasticsearchStore.html) to connect to our elastic cloud deployment. This would help create and index data easily. In the ElasticsearchStore instance, will set embedding to [OpenAIEmbeddings](https://api.python.langchain.com/en/latest/embeddings/langchain.embeddings.openai.OpenAIEmbeddings.html) to embed the texts and elasticsearch index name that will be used in this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "a-t1mglib54F"
   },
   "outputs": [],
   "source": [
    "from langchain_elasticsearch import ElasticsearchStore\n",
    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
    "from getpass import getpass\n",
    "\n",
    "# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#finding-your-cloud-id\n",
    "ELASTIC_CLOUD_ID = getpass(\"Elastic Cloud ID: \")\n",
    "\n",
    "# https://www.elastic.co/search-labs/tutorials/install-elasticsearch/elastic-cloud#creating-an-api-key\n",
    "ELASTIC_API_KEY = getpass(\"Elastic Api Key: \")\n",
    "\n",
    "# https://platform.openai.com/api-keys\n",
    "OPENAI_API_KEY = getpass(\"OpenAI API key: \")\n",
    "\n",
    "embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)\n",
    "\n",
    "vector_store = ElasticsearchStore(\n",
    "    es_cloud_id=ELASTIC_CLOUD_ID,\n",
    "    es_api_key=ELASTIC_API_KEY,\n",
    "    index_name=\"workplace_index\",\n",
    "    embedding=embeddings,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Indexing Data into Elasticsearch\n",
    "\n",
    "Let's download the sample dataset and deserialize the document. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "J8-93TiJsNyK"
   },
   "outputs": [],
   "source": [
    "from urllib.request import urlopen\n",
    "from langchain.llms import OpenAI\n",
    "import json\n",
    "\n",
    "url = \"https://raw.githubusercontent.com/elastic/elasticsearch-labs/main/example-apps/chatbot-rag-app/data/data.json\"\n",
    "\n",
    "response = urlopen(url)\n",
    "data = json.load(response)\n",
    "\n",
    "with open(\"temp.json\", \"w\") as json_file:\n",
    "    json.dump(data, json_file)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "p0cQFDl1b9v4"
   },
   "source": [
    "### Split Documents into Passages\n",
    "\n",
    "We’ll chunk documents into passages in order to improve the retrieval specificity and to ensure that we can provide multiple passages within the context window of the final question answering prompt.\n",
    "\n",
    "Here we are chunking documents into 800 token passages with an overlap of 400 tokens.\n",
    "\n",
    "Here we are using a simple splitter but Langchain offers more advanced splitters to reduce the chance of context being lost."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "dbHEoTF6vBXE"
   },
   "outputs": [],
   "source": [
    "from langchain.document_loaders import JSONLoader\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "\n",
    "def metadata_func(record: dict, metadata: dict) -> dict:\n",
    "    metadata[\"name\"] = record.get(\"name\")\n",
    "    metadata[\"summary\"] = record.get(\"summary\")\n",
    "    metadata[\"url\"] = record.get(\"url\")\n",
    "    metadata[\"category\"] = record.get(\"category\")\n",
    "    metadata[\"updated_at\"] = record.get(\"updated_at\")\n",
    "\n",
    "    return metadata\n",
    "\n",
    "\n",
    "# For more loaders https://python.langchain.com/docs/modules/data_connection/document_loaders/\n",
    "# And 3rd party loaders https://python.langchain.com/docs/modules/data_connection/document_loaders/#third-party-loaders\n",
    "loader = JSONLoader(\n",
    "    file_path=\"temp.json\",\n",
    "    jq_schema=\".[]\",\n",
    "    content_key=\"content\",\n",
    "    metadata_func=metadata_func,\n",
    ")\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n",
    "    chunk_size=512, chunk_overlap=256\n",
    ")\n",
    "docs = loader.load_and_split(text_splitter=text_splitter)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "RmCUl0hxW4lG"
   },
   "source": [
    "### Bulk Import Passages\n",
    "\n",
    "Now that we have split each document into the chunk size of 800, we will now index data to elasticsearch using [ElasticsearchStore.from_documents](https://api.python.langchain.com/en/latest/vectorstores/langchain.vectorstores.elasticsearch.ElasticsearchStore.html#langchain.vectorstores.elasticsearch.ElasticsearchStore.from_documents).\n",
    "\n",
    "We will use Cloud ID,  Password and Index name values set in the `Create cloud deployment` step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = vector_store.from_documents(\n",
    "    docs,\n",
    "    embeddings,\n",
    "    index_name=\"workplace_index\",\n",
    "    es_cloud_id=ELASTIC_CLOUD_ID,\n",
    "    es_api_key=ELASTIC_API_KEY,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "rXJH_MiWejv7"
   },
   "source": [
    "## Asking a question\n",
    "Now that we have the passages stored in Elasticsearch, we can now ask a question to get the relevant passages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OobeBT6rek7Q",
    "outputId": "ba7b3a7a-253e-4e7f-83b9-cec07ebdac09"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---- Answer ----\n",
      "\n",
      "The NASA Sales Team is responsible for understanding the unique market dynamics and cultural nuances of North and South America. It is led by Area Vice-Presidents Laura Martinez (North America) and Gary Johnson (South America), and consists of dedicated account managers, sales representatives, and support staff. The team works to effectively target and engage with customers across the region.\n"
     ]
    }
   ],
   "source": [
    "from langchain.schema.runnable import RunnablePassthrough\n",
    "from langchain.prompts import ChatPromptTemplate\n",
    "from langchain.schema.output_parser import StrOutputParser\n",
    "\n",
    "retriever = vector_store.as_retriever()\n",
    "\n",
    "llm = OpenAI(openai_api_key=OPENAI_API_KEY)\n",
    "\n",
    "ANSWER_PROMPT = ChatPromptTemplate.from_template(\n",
    "    \"\"\"You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Be as verbose and educational in your response as possible. \n",
    "    \n",
    "    context: {context}\n",
    "    Question: \"{question}\"\n",
    "    Answer:\n",
    "    \"\"\"\n",
    ")\n",
    "\n",
    "chain = (\n",
    "    {\"context\": retriever, \"question\": RunnablePassthrough()}\n",
    "    | ANSWER_PROMPT\n",
    "    | llm\n",
    "    | StrOutputParser()\n",
    ")\n",
    "\n",
    "ans = chain.invoke(\"what is the nasa sales team?\")\n",
    "\n",
    "print(\"---- Answer ----\")\n",
    "print(ans)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add Source Tracing\n",
    "RAG can provide clear traceability of the source knowledge used to answer a question. This is important for compliance and regulatory reasons and limiting LLM hallucinations. This is known as source tracking.\n",
    "\n",
    "In this example, we extend the Prompt template to ask the LLM to cite the source of the answer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---- Answer ----\n",
      "The North America South America (NASA) sales team is responsible for serving customers and achieving business objectives across North and South America. The team is led by two Area Vice-Presidents: Laura Martinez is the Area Vice-President of North America, and Gary Johnson is the Area Vice-President of South America. The team consists of dedicated account managers, sales representatives, and support staff. They are responsible for identifying and pursuing new business opportunities, nurturing existing client relationships, and ensuring customer satisfaction.\n",
      "SOURCE: Sales Organization Overview\n"
     ]
    }
   ],
   "source": [
    "from langchain.schema.runnable import RunnablePassthrough\n",
    "from langchain.prompts import ChatPromptTemplate, PromptTemplate\n",
    "from langchain.schema.output_parser import StrOutputParser\n",
    "from langchain.schema import format_document\n",
    "\n",
    "retriever = vector_store.as_retriever()\n",
    "\n",
    "llm = OpenAI(openai_api_key=OPENAI_API_KEY)\n",
    "\n",
    "ANSWER_PROMPT = ChatPromptTemplate.from_template(\n",
    "    \"\"\"\n",
    "You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Be as verbose and educational in your response as possible. \n",
    "Each passage has a SOURCE which is the title of the document. When answering, cite source name of the passages you are answering from below the answer, on a new line, with a prefix of \"SOURCE:\".\n",
    "\n",
    "\n",
    "context: {context}\n",
    "Question: \"{question}\"\n",
    "Answer:\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "DOCUMENT_PROMPT = PromptTemplate.from_template(\n",
    "    \"\"\"\n",
    "---\n",
    "SOURCE: {name}\n",
    "{page_content}\n",
    "---\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "\n",
    "def _combine_documents(\n",
    "    docs, document_prompt=DOCUMENT_PROMPT, document_separator=\"\\n\\n\"\n",
    "):\n",
    "    doc_strings = [format_document(doc, document_prompt) for doc in docs]\n",
    "    return document_separator.join(doc_strings)\n",
    "\n",
    "\n",
    "_context = {\n",
    "    \"context\": retriever | _combine_documents,\n",
    "    \"question\": RunnablePassthrough(),\n",
    "}\n",
    "\n",
    "chain = _context | ANSWER_PROMPT | llm | StrOutputParser()\n",
    "\n",
    "ans = chain.invoke(\"what is the nasa sales team?\")\n",
    "\n",
    "print(\"---- Answer ----\")\n",
    "print(ans)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Returning Passages with Answer\n",
    "\n",
    "In this example, we extend the chain to return the passages back with the answer. This is helpful for the UI to display the source passages, should the user want to read more on the topic. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---- Answer ----\n",
      "The North America South America (NASA) region has two Area Vice-Presidents: Laura Martinez is the Area Vice-President of North America, and Gary Johnson is the Area Vice-President of South America. The NASA sales team consists of dedicated account managers, sales representatives, and support staff, led by their respective Area Vice-Presidents. They are responsible for identifying and pursuing new business opportunities, nurturing existing client relationships, and ensuring customer satisfaction. The teams collaborate closely with other departments, such as marketing, product development, and customer support, to ensure we consistently deliver high-quality products and services to our clients.\n",
      "\n",
      "SOURCE: Sales Organization Overview\n",
      "\n",
      "---- Documents ----\n",
      "Sales Organization Overview\n",
      "Our sales organization is structured to effectively serve our customers and achieve our business objectives across multiple regions. The organization is divided into the following main regions:\n",
      "\n",
      "The Americas: This region includes the United States, Canada, Mexico, as well as Central and South America. The North America South America region (NASA) has two Area Vice-Presidents: Laura Martinez is the Area Vice-President of North America, and Gary Johnson is the Area Vice-President of South America.\n",
      "\n",
      "Europe: Our European sales team covers the entire continent, including the United Kingdom, Germany, France, Spain, Italy, and other countries. The team is responsible for understanding the unique market dynamics and cultural nuances, enabling them to effectively target and engage with customers across the region. The Area Vice-President for Europe is Rajesh Patel.\n",
      "Asia-Pacific: This region encompasses countries such as China, Japan, South Korea, India, Australia, and New Zealand. Our sales team in the Asia-Pacific region works diligently to capitalize on growth opportunities and address the diverse needs of customers in this vast and rapidly evolving market. The Area Vice-President for Asia-Pacific is Mei Li.\n",
      "Middle East & Africa: This region comprises countries across the Middle East and Africa, such as the United Arab Emirates, Saudi Arabia, South Africa, and Nigeria. Our sales team in this region is responsible for navigating the unique market challenges and identifying opportunities to expand our presence and better serve our customers. The Area Vice-President for Middle East & Africa is Jamal Abdi.\n",
      "\n",
      "Each regional sales team consists of dedicated account managers, sales representatives, and support staff, led by their respective Area Vice-Presidents. They are responsible for identifying and pursuing new business opportunities, nurturing existing client relationships, and ensuring customer satisfaction. The teams collaborate closely with other departments, such as marketing, product development, and customer support, to ensure we consistently deliver high-quality products and services to our clients.\n",
      "----\n",
      "Sales Engineering Collaboration\n",
      "As an engineer, it is important to understand the sales team's goals and objectives, as this will help you to provide them with the necessary information, tools, and support to successfully sell your company's products and services.\n",
      "Communication:\n",
      "Effective communication is key to successfully working with the sales team. Make sure to maintain open lines of communication, and be responsive to their questions and concerns. This includes:\n",
      "\n",
      "a. Attending sales meetings and conference calls when required.\n",
      "b. Providing regular product updates and training sessions to the sales team.\n",
      "c. Being available to answer technical questions and clarifications.\n",
      "Collaboration:\n",
      "Collaborate with the sales team in developing and refining sales materials, such as product presentations, demos, and technical documents. This will ensure that the sales team has accurate and up-to-date information to present to clients.\n",
      "\n",
      "Additionally, work closely with the sales team on customer projects or product customizations, providing technical guidance, and ensuring that the solutions meet the customer's requirements.\n",
      "Customer Engagement:\n",
      "At times, engineers may be asked to join sales meetings or calls with potential clients to provide technical expertise. In these situations, it is important to:\n",
      "\n",
      "a. Be prepared and understand the customer's needs and pain points.\n",
      "b. Clearly explain the technical aspects of the product or solution in a simple language that the customer can understand.\n",
      "c. Address any concerns or questions the customer may have.\n",
      "Continuous Improvement:\n",
      "Actively seek feedback from the sales team regarding product performance, customer experiences, and market trends. Use this feedback to identify areas of improvement and collaborate with other engineers to enhance the product or service offerings.\n",
      "Mutual Respect and Support:\n",
      "It is essential to treat your colleagues in the sales team with respect and professionalism. Recognize and appreciate their efforts in promoting and selling the company's products and services. In turn, the sales team should also respect and appreciate the technical expertise and knowledge of the engineering team.\n",
      "\n",
      "By working together, both the engineering and sales teams can contribute to the overall success of the company.\n",
      "\n",
      "Conclusion:\n",
      "Collaboration between engineers and the sales team is crucial for a tech company's success. By understanding each other's roles, maintaining effective communication, collaborating on projects, and supporting one another, both teams can work together to achieve the company's goals and ensure customer satisfaction.\n",
      "----\n",
      "Fy2024 Company Sales Strategy\n",
      "III. Action Plans\n",
      "A. Sales Team Development:\n",
      "Expand the sales team to cover new markets and industries.\n",
      "Provide ongoing training to sales staff on product knowledge, sales techniques, and industry trends.\n",
      "Implement a performance-based incentive system to reward top performers.\n",
      "\n",
      "B. Marketing and Promotion:\n",
      "Develop targeted marketing campaigns for different customer segments and industries.\n",
      "Leverage digital marketing channels to increase brand visibility and lead generation.\n",
      "Participate in industry events and trade shows to showcase our products and services.\n",
      "\n",
      "C. Partner Ecosystem:\n",
      "Strengthen existing partnerships and establish new strategic alliances to expand market reach.\n",
      "Collaborate with partners on joint marketing and sales initiatives.\n",
      "Provide partner training and support to ensure they effectively represent our products and services.\n",
      "\n",
      "D. Customer Success:\n",
      "Implement a proactive customer success program to improve customer retention and satisfaction.\n",
      "Develop a dedicated customer support team to address customer inquiries and concerns promptly.\n",
      "Collect and analyze customer feedback to identify areas for improvement in our products, services, and processes.\n",
      "\n",
      "IV. Monitoring and Evaluation\n",
      "Establish key performance indicators (KPIs) to track progress toward our objectives.\n",
      "Conduct regular sales team meetings to review performance, share best practices, and address challenges.\n",
      "Conduct quarterly reviews of our sales strategy to ensure alignment with market trends and adjust as needed.\n",
      "\n",
      "By following this sales strategy for fiscal year 2024, our tech company aims to achieve significant growth and success in our target markets, while also providing exceptional value and service to our customers.\n",
      "----\n",
      "Sales Engineering Collaboration\n",
      "Title: Working with the Sales Team as an Engineer in a Tech Company\n",
      "\n",
      "Introduction:\n",
      "As an engineer in a tech company, collaboration with the sales team is essential to ensure the success of the company's products and services. This guidance document aims to provide an overview of how engineers can effectively work with the sales team, fostering a positive and productive working environment.\n",
      "Understanding the Sales Team's Role:\n",
      "The sales team is responsible for promoting and selling the company's products and services to potential clients. Their role involves establishing relationships with customers, understanding their needs, and ensuring that the offered solutions align with their requirements.\n",
      "\n",
      "As an engineer, it is important to understand the sales team's goals and objectives, as this will help you to provide them with the necessary information, tools, and support to successfully sell your company's products and services.\n",
      "Communication:\n",
      "Effective communication is key to successfully working with the sales team. Make sure to maintain open lines of communication, and be responsive to their questions and concerns. This includes:\n",
      "\n",
      "a. Attending sales meetings and conference calls when required.\n",
      "b. Providing regular product updates and training sessions to the sales team.\n",
      "c. Being available to answer technical questions and clarifications.\n",
      "Collaboration:\n",
      "Collaborate with the sales team in developing and refining sales materials, such as product presentations, demos, and technical documents. This will ensure that the sales team has accurate and up-to-date information to present to clients.\n",
      "\n",
      "Additionally, work closely with the sales team on customer projects or product customizations, providing technical guidance, and ensuring that the solutions meet the customer's requirements.\n",
      "Customer Engagement:\n",
      "At times, engineers may be asked to join sales meetings or calls with potential clients to provide technical expertise. In these situations, it is important to:\n",
      "----\n"
     ]
    }
   ],
   "source": [
    "from langchain.schema.runnable import RunnableMap\n",
    "from langchain.prompts import ChatPromptTemplate, PromptTemplate\n",
    "from langchain.schema import format_document\n",
    "from operator import itemgetter\n",
    "\n",
    "retriever = vector_store.as_retriever()\n",
    "\n",
    "llm = OpenAI(openai_api_key=OPENAI_API_KEY)\n",
    "\n",
    "ANSWER_PROMPT = ChatPromptTemplate.from_template(\n",
    "    \"\"\"\n",
    "You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Be as verbose and educational in your response as possible. \n",
    "Each passage has a SOURCE which is the title of the document. When answering, cite source name of the passages you are answering from below the answer, on a new line, with a prefix of \"SOURCE:\".\n",
    "\n",
    "context: {context}\n",
    "Question: {question}\n",
    "Answer:\n",
    "\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "DOCUMENT_PROMPT = PromptTemplate.from_template(\n",
    "    \"\"\"\n",
    "---\n",
    "SOURCE: {name}\n",
    "{page_content}\n",
    "---\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "\n",
    "def _combine_documents(\n",
    "    docs, document_prompt=DOCUMENT_PROMPT, document_separator=\"\\n\\n\"\n",
    "):\n",
    "    doc_strings = [format_document(doc, document_prompt) for doc in docs]\n",
    "    return document_separator.join(doc_strings)\n",
    "\n",
    "\n",
    "retrieved_documents = RunnableMap(\n",
    "    docs=itemgetter(\"question\") | retriever,\n",
    "    question=itemgetter(\"question\"),\n",
    ")\n",
    "\n",
    "_context = {\n",
    "    \"context\": lambda x: _combine_documents(x[\"docs\"]),\n",
    "    \"question\": lambda x: x[\"question\"],\n",
    "}\n",
    "\n",
    "answer = {\n",
    "    \"answer\": _context | ANSWER_PROMPT | llm,\n",
    "    \"docs\": itemgetter(\"docs\"),\n",
    "}\n",
    "\n",
    "chain = retrieved_documents | answer\n",
    "\n",
    "ans = chain.invoke({\"question\": \"what is the nasa sales team?\"})\n",
    "\n",
    "print(\"---- Answer ----\")\n",
    "print(ans[\"answer\"])\n",
    "print()\n",
    "print(\"---- Documents ----\")\n",
    "for doc in ans[\"docs\"]:\n",
    "    print(doc.metadata[\"name\"])\n",
    "    print(doc.page_content)\n",
    "    print(\"----\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conversational Question Answering\n",
    "We have achieved getting answers to questions, but what if we want to ask follow up questions? We can use the answer from the previous question as the context for the next question. This is known as conversational question answering.\n",
    "\n",
    "In this example, we extend the chain to use the answer from the previous question as the context for the next question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---- Answer ----\n",
      "The objectives for fiscal year 2024 are to increase revenue by 20% compared to fiscal year 2023, expand market share in key segments by 15%, retain 95% of existing customers and increase customer satisfaction ratings, and launch at least two new products or services in high-demand market segments. SOURCE: Fy2024 Company Sales Strategy\n"
     ]
    }
   ],
   "source": [
    "from langchain.schema.runnable import RunnableMap\n",
    "from langchain.prompts import ChatPromptTemplate, PromptTemplate\n",
    "from langchain.schema import format_document\n",
    "from operator import itemgetter\n",
    "\n",
    "retriever = vector_store.as_retriever()\n",
    "\n",
    "llm = OpenAI(openai_api_key=OPENAI_API_KEY)\n",
    "\n",
    "ANSWER_PROMPT = ChatPromptTemplate.from_template(\n",
    "    \"\"\"\n",
    "You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Be as verbose and educational in your response as possible. \n",
    "Each passage has a SOURCE which is the title of the document. When answering, cite source name of the passages you are answering from below the answer, on a new line, with a prefix of \"SOURCE:\".\n",
    "\n",
    "context: \n",
    "{context}\n",
    "\n",
    "Question: {question}\n",
    "Answer:\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "DOCUMENT_PROMPT = PromptTemplate.from_template(\n",
    "    \"\"\"\n",
    "---\n",
    "SOURCE: {name}\n",
    "{page_content}\n",
    "---\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(\n",
    "    \"\"\"Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question, in its original language.\n",
    "\n",
    "Chat History:\n",
    "{chat_history}\n",
    "Follow Up Input: {question}\n",
    "\"\"\"\n",
    ")\n",
    "\n",
    "standalone_question = RunnableMap(\n",
    "    standalone_question=RunnablePassthrough.assign(\n",
    "        chat_history=lambda x: _format_chat_history(x[\"chat_history\"])\n",
    "    )\n",
    "    | CONDENSE_QUESTION_PROMPT\n",
    "    | llm\n",
    "    | StrOutputParser(),\n",
    ")\n",
    "\n",
    "\n",
    "def _format_chat_history(chat_history) -> str:\n",
    "    buffer = \"\"\n",
    "    for dialogue_turn in chat_history:\n",
    "        human = \"Human: \" + dialogue_turn[0]\n",
    "        ai = \"Assistant: \" + dialogue_turn[1]\n",
    "        buffer += \"\\n\" + \"\\n\".join([human, ai])\n",
    "    return buffer\n",
    "\n",
    "\n",
    "def _combine_documents(\n",
    "    docs, document_prompt=DOCUMENT_PROMPT, document_separator=\"\\n\\n\"\n",
    "):\n",
    "    doc_strings = [format_document(doc, document_prompt) for doc in docs]\n",
    "    return document_separator.join(doc_strings)\n",
    "\n",
    "\n",
    "retrieved_documents = RunnableMap(\n",
    "    docs=itemgetter(\"standalone_question\") | retriever,\n",
    "    question=itemgetter(\"standalone_question\"),\n",
    ")\n",
    "\n",
    "_context = {\n",
    "    \"context\": lambda x: _combine_documents(x[\"docs\"]),\n",
    "    \"question\": lambda x: x[\"question\"],\n",
    "}\n",
    "\n",
    "answer = {\n",
    "    \"answer\": _context | ANSWER_PROMPT | llm,\n",
    "    \"docs\": itemgetter(\"docs\"),\n",
    "}\n",
    "\n",
    "chain = standalone_question | retrieved_documents | answer\n",
    "\n",
    "ans = chain.invoke(\n",
    "    {\n",
    "        \"question\": \"What are their objectives?\",\n",
    "        \"chat_history\": [\n",
    "            \"What is the nasa sales team?\",\n",
    "            \"The sales team of NASA consists of Laura Martinez, the Area \"\n",
    "            \"Vice-President of North America, and Gary Johnson, the Area \"\n",
    "            \"Vice-President of South America.\"\n",
    "            \"SOURCE: Sales Organization Overview\",\n",
    "        ],\n",
    "    }\n",
    ")\n",
    "\n",
    "print(\"---- Answer ----\")\n",
    "print(ans[\"answer\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Next Steps\n",
    "We have shown how to use Langchain to build a question answering system. We have shown how to index data into Elasticsearch, ask a question and use the answer from the previous question as the context for the next question."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.11.3 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.3"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
